#include "pch.h"
#include "LoopBuffer.h"
#include "WebSocket.h"
#include <map>

using std::map;
using std::pair;

#pragma region "Session"
#define MEMORY_NO_ENJOY() {	\
							RELEASE_SOCKET(socket); \
							varbuf=_buf; \
							recvBuffer.buf =_buf; \
							recvBuffer.len = SESSION_FIX_RECV_BUFFER_SIZE - 1; \
							postRecv(); \
						  }

typedef map<const char*, const char*, keycmp> HttpHead;
const char* magicStr = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

LoopBuffer Session::recvBufferPool(RECV_BUFFER_POOL_SIZE);
LoopBuffer Session::sendBufferPool(SEND_BUFFER_POOL_SIZE);

static void httpAnalysis(char* buf, HttpHead& http, size_t len) {
	char* n, * v;
	n = v = buf;
	bool bi = false;
	for (size_t i = 0; i < len; i++) {
		switch (buf[i])
		{
		case ':':
			buf[i] = 0;
			v = &buf[i + 1];
			if ((*v) == ' ') v++;
			continue;
		case '\r':
		case '\n':
			buf[i] = 0;
			if (bi) {
				http.insert(pair<char*, char*>(n, v));
				n = &buf[i + 1];
			}
			bi = true;
			continue;
		default:
			if (bi) {
				n = &buf[i];
				bi = false;
			}
			continue;
		}

	}
};
static void msgDecode(char* buf, size_t len, char* msk) {
	size_t len2 = len / 4;
	size_t len3 = len % 4;
	int* p = (int*)buf;
	int* m = (int*)msk;
	for (size_t i = 0; i < len2; i++) {
		*p++ ^= *m;
	}
	buf = (char*)p;
	for (size_t i = 0; i < len3; i++) {
		*buf++ ^= msk[i % 4];
	}
}

BOOL Session::handShake(char* msg, size_t len)
{
	HttpHead httpHeader;
	httpAnalysis(msg, httpHeader, len);
	const char* Sec = "Sec-WebSocket-Key";
	const char* http = "GET / HTTP/1.1";
	if (!httpHeader[http])return false;


	const char* p = httpHeader[Sec];
	if (!p) return false;

	strcat_s((char*)p, 1024 - (p - msg), magicStr);

	unsigned char has[20];
	char bash[40];
	SHA1_String((const unsigned char*)p, strlen(p), has);
	BASE64_Encode((const BYTE*)has, 20, bash);
	msg[0] = 0;
	const char* response = "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: %s\r\n\r\n";
	len = sprintf_s(msg, SESSION_FIX_RECV_BUFFER_SIZE, response, bash);

	OverLapped ovl;
	ZeroMemory(&ovl, sizeof(ovl));
	ovl.sendBuf.buf = msg;
	ovl.sendBuf.len = len;
	ovl.session = this;
	ovl.ioType = IO_TYPE::HANDSHAKE;
	if (WSASend(socket, &(ovl.sendBuf), 1, NULL, 0, &ovl, NULL) != NO_ERROR) {
		if (WSA_IO_PENDING != WSAGetLastError()) {
			return false;
		}
	}
	return true;
}
BOOL Session::RecvMessage(char* buf, size_t len)
{
	PMsgHead ph = (PMsgHead)buf;
	size_t l = ph->getPlayLoad();
	size_t s = ph->getHeadSize();
	buf[len] = 0;
	char* c = buf;
	switch (ph->opCode) {
	case 0:
		msgLen += l;
		break;
	case 1:
		msgLen = l;
		break;
	default:
		if (varbuf != _buf) {
			recvBufferPool.free(varbuf);
		}
		return true;
	}
	msgDecode(buf + s, l, ph->getSHAMask());
	if (varbuf == _buf) {
		c = recvBufferPool.get(this, len + 1);
		if (c) {
			CopyMemory(c, buf, len);
		}
		else {
			return false;
		}
	}
	packetList.push_back(c);

	if (ph->fin) {
		c = recvBufferPool.get(this, msgLen + 1);
		if (!c) return false;
		char* p = c;
		for (vector<char*>::iterator it = packetList.begin(); it != packetList.end(); it++) {
			size_t l, s;
			ph = (PMsgHead)(*it);
			l = ph->getPlayLoad();
			s = ph->getHeadSize();
			CopyMemory(c, (*it) + s, l);
			c += l;
			recvBufferPool.free(*it);
		}
		packetList.clear();
		*c = 0;
		ws->onMessage(*this, p, msgLen);
	}
	return true;
}
void Session::Recv(size_t len)
{
	if (CONNECT_STATUS::pending == status) {
		if (handShake(_buf, len)) {
			status = CONNECT_STATUS::connected;
		}
		postRecv();
	}
	else {
		char* buf = varbuf;
		PMsgHead ph = (PMsgHead)buf;
		position += len;
		size_t msglen = ph->getLength(position);
		while (msglen && msglen <= position) {
			if (!RecvMessage(buf, msglen)) {
				MEMORY_NO_ENJOY();
				return;
			};
			buf += msglen;
			position -= msglen;
			ph = (PMsgHead)buf;
			msglen = ph->getLength(position);
		}
		if (msglen) {
			if (varbuf != _buf) {
				recvBuffer.buf += len;
				recvBuffer.len -= len;
			}
			else {
				varbuf = recvBufferPool.get(this, msglen + 1);
				if (!varbuf) {
					MEMORY_NO_ENJOY()
						return;
				}
				CopyMemory(varbuf, buf, position);
				recvBuffer.buf = varbuf + position;
				recvBuffer.len = msglen - position;
			}
		}
		else {
			CopyMemory(_buf, buf, position);
			varbuf = _buf;
			recvBuffer.buf = _buf + position;
			recvBuffer.len = SESSION_FIX_RECV_BUFFER_SIZE - position - 1;
		}
		postRecv();
	}
}

bool Session::SendMsg(char* buf, size_t len)
{
	PMsgHead ph;
	POverLapped ovl;
	size_t szHead;
	szHead = ph->getSendHeadSize(len);
	char* p;
	p = sendBufferPool.get(this, szHead + len + sizeof(OverLapped) + 16);
	while (p == NULL) {
		Sleep(100);
	};
	ph = (PMsgHead)p;
	ph->mkSendBuf(buf, len);

	ovl = (POverLapped)(((DWORD)p + szHead + len + 8) & 0xfffffff8);
	ZeroMemory(ovl, sizeof(OVERLAPPED));

	ovl->session = this;
	ovl->ioType = IO_TYPE::SEND;
	ovl->sendBuf.len = szHead + len;
	ovl->sendBuf.buf = p;
	DWORD dwflag = 0;
	if (WSASend(socket, &(ovl->sendBuf), 1, NULL, dwflag, ovl, NULL) != NO_ERROR) {
		if (WSA_IO_PENDING != WSAGetLastError()) {
			sendBufferPool.free(p);
			return false;
		}
	}
	return true;
}

inline void Session::postRecv()
{
	DWORD dwFlags = 0;
	ZeroMemory(&recvOvl, sizeof(WSAOVERLAPPED));
	int nBytesRecv = WSARecv(socket, &recvBuffer, 1, NULL, &dwFlags, &recvOvl, NULL);

	if ((SOCKET_ERROR == nBytesRecv) && (WSA_IO_PENDING != WSAGetLastError()))
	{
		debug("From Session postRec sessioncnt:%d\r\n", ws->sessionCNT);
		ws->recyleSession(this);

	}
}
#pragma endregion

